import time
from typing import List
import torch
from rich import print
from transformers import AutoTokenizer, AutoModelForMaskedLM

from utils.verbalizer import Verbalizer
from data_handle.data_preprocess import convert_examples
from utils.common_utils import convert_logits_to_ids
from p_tuning_config import *

pc = p_tuning_config()
device = pc.device

model_path = 'D:\Project\AIStudent\LLMProject\BERT_Fintuning\PET\checkpoints\model_best'
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForMaskedLM.from_pretrained(model_path)
model.to(device)
model.eval()

verbalizer = Verbalizer(verbalizer_file=pc.verbalizer,
                        tokenizer=tokenizer,
                        max_label_len=pc.max_label_len)

def inference(contents: List[str]):
    with torch.no_grad():
        start_time = time.time()
        examples = {'text' : contents}

        tokenizer_output = convert_examples(
            examples,
            tokenizer,
            max_label_len=pc.max_label_len,
            max_seq_length=128,
            train_mode=False,
            return_tensor=True
        )

        logits = model(input_ids=tokenizer_output['input_ids'].to(device),
                       attention_mask=tokenizer_output['attention_mask'].to(device),
                       token_type_ids=tokenizer_output['token_type_ids'].to(device),).logits

        predictions = convert_logits_to_ids(logits, tokenizer_output['mask_positions']).cpu().numpy().tolist()
        predictions = verbalizer.batch_find_main_label(predictions)
        predictions = [ele['label'] for ele in predictions]
        end_time = time.time()
        print(f'inference time: {end_time - start_time}')
        return predictions

if __name__ == '__main__':
    contents = [
        '只能呵呵哒了，电视太旧，没有电视遥控器！被子潮潮的，像没有干一样！（最近都是晴天的情况下）插座只有电视下面有，床头根本没有！一点四星级的档次都没有，还不如三星的格林豪泰！',
        '东西很好，质量也不错，快递很快，服务态度也很好，我非常满意！'
        ]
    print("预测结果")
    result = inference(contents)
    print(result)
    new_dict = {}
    for i in range(len(contents)):
        new_dict[contents[i]] = result[i]

    print(new_dict)